package com.tianji.aigc.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.IdUtil;
import com.tianji.aigc.constant.ToolConstant;
import com.tianji.aigc.dto.ChatDTO;
import com.tianji.aigc.entity.ChatSession;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.properties.SessionProperties;
import com.tianji.aigc.service.ChatService;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.aigc.tools.result.ToolResultHolder;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.exceptions.BadRequestException;
import com.tianji.common.utils.UserContext;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 对话业务实现类
 *
 * @Author mr.wu
 * @Date 2025-8-17 15:28
 */
@Service
public class ChatServiceImpl implements ChatService {

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private IChatSessionService chatSessionService;

    //存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全，如果分布式环境可考虑用redis
    private static Map<String, Boolean> SESSION_MAP = new ConcurrentHashMap<String, Boolean>();

    @Autowired
    private ChatMemory chatMemory;

    @Autowired
    private VectorStore vectorStore;
    @Autowired
    private SessionProperties sessionProperties;

    /**
     * 聊天对话
     *
     * @param dto
     * @return
     */
    @Override
    public Flux<ChatEventVO> chat(ChatDTO dto) {
        //得到sessionId
        String sessionId = dto.getSessionId();
        //判断是会话ID是否存在
        ChatSession chatSession = chatSessionService.lambdaQuery().eq(ChatSession::getSessionId, dto.getSessionId()).one();
        if (chatSession == null) {
            throw new BadRequestException("会话不存在，请先创建会话~");
        }
        //更新第一次的会话标题
        chatSessionService.updateSessionTitle(dto.getQuestion(),dto.getSessionId(),UserContext.getUser());

        //保存中断的会话内容
        StringBuffer assistantContent = new StringBuffer();
        //对话ID
        String converationId = UserContext.getUser() + ":" + dto.getSessionId();


        //生成唯一请求ID
        String requestId = IdUtil.fastSimpleUUID();

        ChatEventVO STOP_EVENT = ChatEventVO.builder()  // 标记输出结束
                .eventType(ChatEventTypeEnum.STOP.getValue())
                .build();

        //调用大模型进行对话
        return chatClient.prompt()
                .user(dto.getQuestion())
                .advisors(x -> x
                        //RAG查询，由于数据量少，直接查询全部
                        .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.builder().query("").topK(999).build()))
                        .param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, converationId))//保存会话ID，用来区分不同的会话
                .toolContext(MapUtil.<String, Object>builder()
                        .put(ToolConstant.REQUEST_ID, requestId)//保存请求ID到模型对象中的一个map集合中
                        .put(ToolConstant.USER_ID, UserContext.getUser())//保存用户ID到模型对象中的一个map集合中
                        .build()) //将请求ID传递给课程Tool
                .stream()
//                .content()
                .chatResponse()
                .doFirst(() -> SESSION_MAP.put(sessionId, true))//标记开始输出
                .doOnComplete(() -> SESSION_MAP.remove(sessionId))//完成之后删除标记
                .doOnError(throwable -> SESSION_MAP.remove(sessionId))//输出异常时，清除标记
                .doOnCancel(() -> this.saveRedis(converationId, assistantContent.toString()))
                .takeWhile(takewile -> SESSION_MAP.getOrDefault(sessionId, false))//判断是否继续输出
                .map(x -> {
                    //先判断是不是结束的内容
                    String finishReason = x.getResult().getMetadata().getFinishReason();
                    if (finishReason.equals("STOP")) {
                        //得到消息ID
                        String messageId = x.getMetadata().getId();
                        //把消息Id存储到ToolResultHolder中
                        ToolResultHolder.put(messageId,ToolConstant.REQUEST_ID,requestId);
                    }
                    //存储对话内容
                    String text = x.getResult().getOutput().getText();
                    assistantContent.append(text);//追加对话信息
                    ChatEventVO chatEventVO = ChatEventVO.builder().
                            eventData(text)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                    return chatEventVO;
                })
                .concatWith(Flux.defer(() -> {
                    //流式数据返回前进行判断，从结果保持器中获取课程Tool存储的数据
                    Map<String, Object> resultMap = ToolResultHolder.get(requestId);
                    if (CollUtil.isNotEmpty(resultMap)) {
                        //如果数据非空，说明课程Tool被调用有结果了，将这个结果封装到结果对象中
                        ChatEventVO chatEventVO = ChatEventVO.builder().eventData(resultMap)
                                .eventType(ChatEventTypeEnum.PARAM.getValue()).build();

                        //清除请求ID对应的存储的数据（用完就清除）
                        ToolResultHolder.remove(requestId);
                        return Flux.just(chatEventVO, STOP_EVENT);
                    }
                    //输出结束
                    return Flux.just(STOP_EVENT);
                }));
    }

    /**
     * 保存对话中断的信息到redis缓存中
     *
     * @param converationId 对话ID
     * @param content       会话内容
     */
    private void saveRedis(String converationId, String content) {
        chatMemory.add(converationId, new AssistantMessage(content));
    }

    /**
     * 停止对话生成
     *
     * @param sessionId
     */
    @Override
    public void stop(String sessionId) {
        SESSION_MAP.put(sessionId, true);//设置标记
    }

    /**
     * 文本聊天
     * @param question
     * @return
     */
    @Override
    public String chatText(String question) {
        return chatClient
                .prompt()
                .system(x->x.text(sessionProperties.getText()))
                .user(question)
                .call()
                .content();
    }
}